"""
重试机制模块

提供灵活的重试策略和装饰器。
"""

import time
import random
from typing import Callable, Optional, Type, Union, Tuple, Any
from functools import wraps
from .exceptions import ESArchiveError, is_retryable_error, RetryableError
from .logger import get_module_logger

logger = get_module_logger(__name__)


class RetryStrategy:
    """重试策略基类"""
    
    def get_delay(self, attempt: int) -> float:
        """获取重试延迟时间
        
        Args:
            attempt: 重试次数 (从1开始)
            
        Returns:
            延迟时间(秒)
        """
        raise NotImplementedError
    
    def should_retry(self, attempt: int, max_attempts: int, error: Exception) -> bool:
        """判断是否应该重试
        
        Args:
            attempt: 当前重试次数
            max_attempts: 最大重试次数
            error: 异常对象
            
        Returns:
            是否应该重试
        """
        return attempt < max_attempts and is_retryable_error(error)


class FixedDelayStrategy(RetryStrategy):
    """固定延迟重试策略"""
    
    def __init__(self, delay: float = 1.0):
        """初始化
        
        Args:
            delay: 固定延迟时间(秒)
        """
        self.delay = delay
    
    def get_delay(self, attempt: int) -> float:
        return self.delay


class ExponentialBackoffStrategy(RetryStrategy):
    """指数退避重试策略"""
    
    def __init__(
        self,
        initial_delay: float = 1.0,
        max_delay: float = 60.0,
        multiplier: float = 2.0,
        jitter: bool = True
    ):
        """初始化
        
        Args:
            initial_delay: 初始延迟时间(秒)
            max_delay: 最大延迟时间(秒)
            multiplier: 延迟倍数
            jitter: 是否添加随机抖动
        """
        self.initial_delay = initial_delay
        self.max_delay = max_delay
        self.multiplier = multiplier
        self.jitter = jitter
    
    def get_delay(self, attempt: int) -> float:
        delay = self.initial_delay * (self.multiplier ** (attempt - 1))
        delay = min(delay, self.max_delay)
        
        if self.jitter:
            # 添加±25%的随机抖动
            jitter_range = delay * 0.25
            delay += random.uniform(-jitter_range, jitter_range)
        
        return max(0, delay)


class LinearBackoffStrategy(RetryStrategy):
    """线性退避重试策略"""
    
    def __init__(self, initial_delay: float = 1.0, increment: float = 1.0, max_delay: float = 30.0):
        """初始化
        
        Args:
            initial_delay: 初始延迟时间(秒)
            increment: 每次增加的延迟时间(秒)
            max_delay: 最大延迟时间(秒)
        """
        self.initial_delay = initial_delay
        self.increment = increment
        self.max_delay = max_delay
    
    def get_delay(self, attempt: int) -> float:
        delay = self.initial_delay + (attempt - 1) * self.increment
        return min(delay, self.max_delay)


class RetryManager:
    """重试管理器"""
    
    def __init__(
        self,
        max_attempts: int = 3,
        strategy: Optional[RetryStrategy] = None,
        exceptions: Optional[Tuple[Type[Exception], ...]] = None
    ):
        """初始化重试管理器
        
        Args:
            max_attempts: 最大重试次数
            strategy: 重试策略
            exceptions: 需要重试的异常类型
        """
        self.max_attempts = max_attempts
        self.strategy = strategy or ExponentialBackoffStrategy()
        self.exceptions = exceptions or (Exception,)
    
    def execute(self, func: Callable, *args, **kwargs) -> Any:
        """执行函数并处理重试
        
        Args:
            func: 要执行的函数
            *args: 函数参数
            **kwargs: 函数关键字参数
            
        Returns:
            函数执行结果
            
        Raises:
            最后一次执行的异常
        """
        last_exception = None
        
        for attempt in range(1, self.max_attempts + 1):
            try:
                logger.debug(f"执行函数 {func.__name__}, 尝试 {attempt}/{self.max_attempts}")
                result = func(*args, **kwargs)
                
                if attempt > 1:
                    logger.info(f"函数 {func.__name__} 在第 {attempt} 次尝试后成功")
                
                return result
                
            except self.exceptions as e:
                last_exception = e
                
                if not self.strategy.should_retry(attempt, self.max_attempts, e):
                    logger.error(f"函数 {func.__name__} 执行失败，不可重试: {e}")
                    raise
                
                if attempt < self.max_attempts:
                    delay = self.strategy.get_delay(attempt)
                    logger.warning(
                        f"函数 {func.__name__} 执行失败 (尝试 {attempt}/{self.max_attempts}), "
                        f"将在 {delay:.2f}s 后重试: {e}"
                    )
                    time.sleep(delay)
                else:
                    logger.error(f"函数 {func.__name__} 执行失败，已达到最大重试次数: {e}")
        
        # 如果所有重试都失败了，抛出最后一个异常
        raise last_exception


def retry(
    max_attempts: int = 3,
    strategy: Optional[RetryStrategy] = None,
    exceptions: Optional[Tuple[Type[Exception], ...]] = None
):
    """重试装饰器
    
    Args:
        max_attempts: 最大重试次数
        strategy: 重试策略
        exceptions: 需要重试的异常类型
    """
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            retry_manager = RetryManager(max_attempts, strategy, exceptions)
            return retry_manager.execute(func, *args, **kwargs)
        return wrapper
    return decorator


def retry_on_es_error(max_attempts: int = 3, strategy: Optional[RetryStrategy] = None):
    """ES错误重试装饰器
    
    Args:
        max_attempts: 最大重试次数
        strategy: 重试策略
    """
    import elasticsearch
    
    es_exceptions = (
        elasticsearch.ConnectionError,
        elasticsearch.ConnectionTimeout,
        elasticsearch.TransportError,
    )
    
    return retry(max_attempts, strategy, es_exceptions)


def retry_on_storage_error(max_attempts: int = 3, strategy: Optional[RetryStrategy] = None):
    """存储错误重试装饰器
    
    Args:
        max_attempts: 最大重试次数
        strategy: 重试策略
    """
    storage_exceptions = (
        IOError,
        OSError,
    )
    
    return retry(max_attempts, strategy, storage_exceptions)


class CircuitBreaker:
    """熔断器
    
    当错误率过高时，暂时停止调用以避免系统过载。
    """
    
    def __init__(
        self,
        failure_threshold: int = 5,
        recovery_timeout: float = 60.0,
        expected_exception: Type[Exception] = Exception
    ):
        """初始化熔断器
        
        Args:
            failure_threshold: 失败阈值
            recovery_timeout: 恢复超时时间(秒)
            expected_exception: 预期的异常类型
        """
        self.failure_threshold = failure_threshold
        self.recovery_timeout = recovery_timeout
        self.expected_exception = expected_exception
        
        self.failure_count = 0
        self.last_failure_time = None
        self.state = "CLOSED"  # CLOSED, OPEN, HALF_OPEN
    
    def call(self, func: Callable, *args, **kwargs) -> Any:
        """调用函数并处理熔断
        
        Args:
            func: 要调用的函数
            *args: 函数参数
            **kwargs: 函数关键字参数
            
        Returns:
            函数执行结果
        """
        if self.state == "OPEN":
            if self._should_attempt_reset():
                self.state = "HALF_OPEN"
                logger.info("熔断器进入半开状态，尝试恢复")
            else:
                raise ESArchiveError("熔断器开启，拒绝调用")
        
        try:
            result = func(*args, **kwargs)
            self._on_success()
            return result
        except self.expected_exception as e:
            self._on_failure()
            raise
    
    def _should_attempt_reset(self) -> bool:
        """检查是否应该尝试重置熔断器"""
        return (
            self.last_failure_time is not None and
            time.time() - self.last_failure_time >= self.recovery_timeout
        )
    
    def _on_success(self) -> None:
        """成功时的处理"""
        if self.state == "HALF_OPEN":
            self.state = "CLOSED"
            logger.info("熔断器恢复到关闭状态")
        
        self.failure_count = 0
    
    def _on_failure(self) -> None:
        """失败时的处理"""
        self.failure_count += 1
        self.last_failure_time = time.time()
        
        if self.failure_count >= self.failure_threshold:
            self.state = "OPEN"
            logger.warning(f"熔断器开启，失败次数: {self.failure_count}")


def circuit_breaker(
    failure_threshold: int = 5,
    recovery_timeout: float = 60.0,
    expected_exception: Type[Exception] = Exception
):
    """熔断器装饰器
    
    Args:
        failure_threshold: 失败阈值
        recovery_timeout: 恢复超时时间(秒)
        expected_exception: 预期的异常类型
    """
    breaker = CircuitBreaker(failure_threshold, recovery_timeout, expected_exception)
    
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            return breaker.call(func, *args, **kwargs)
        return wrapper
    return decorator
